import requests
import json
import time

def ollama_qwen3_answer(prompt, model="qwen3:4b", base_url="http://localhost:11434"):
    # 构建API请求地址
    url = f"{base_url}/api/generate"

    # 准备请求数据
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False  # 非流式响应
    }

    try:
        # 发送POST请求
        response = requests.post(
            url,
            headers={"Content-Type": "application/json"},
            data=json.dumps(payload),
            timeout=600  # 设置超时时间
        )

        # 检查请求是否成功
        response.raise_for_status()

        # 解析响应内容
        result = response.json()
        return result.get("response", "未获取到有效回答")

    except requests.exceptions.RequestException as e:
        return f"请求出错: {str(e)}"
    except json.JSONDecodeError:
        return "响应解析失败，无法识别的JSON格式"


# 示例使用
if __name__ == "__main__":
    # 确保已通过命令下载模型: ollama pull qwen3
    question = "使用“distinctive”生一个英文句子，英文单词数大于5个字，不超过20"

    # 调用模型
    start_time = time.time()
    answer = ollama_qwen3_answer(question)
    end_time = time.time()
    execution_time = end_time - start_time
    print(f"函数执行时间: {execution_time:.6f} 秒")
    # 打印结果
    start = answer.index("</think>")
    if start != -1:
        answer = answer[start + 8:].strip()

    print(f"{answer}")
